//
// Created by hujin on 2023/3/8.
//

#include <cosi/objectives.h>
#include <cmath>


template <typename T> T Objectives<T>::merge(NodeStatL2<T> & node1, NodeStatL2<T> & node2, T e) {
    return  (node1.volume - node1.p_o) * log2(node1.volume)
          + (node2.volume - node2.p_o) * log2(node2.volume)
          - (node1.volume + node2.volume - node1.p_o - node2.p_o + e) * log2(node1.volume + node2.volume)
          + e * log2(node1.volume_h);
}


template<typename T>
T Objectives<T>::merge(NodeStatL3<T> & node1, NodeStatL3<T> & node2, T e) {
    return (node1.p_a - node1.p_o) * log2(node1.volume)
           + (node2.p_a - node2.p_o) * log2(node2.volume)
           - (node1.p_a + node2.p_a - node1.p_o - node2.p_o + e) * log2(node1.volume + node2.volume)
           + e * log2(node1.volume_h);
}


template<typename T>
T Objectives<T>::lift(NodeStatL2<T> & s, NodeStatL2<T> & leaf, T e) {
    NodeStatL2<T> s_old = {s.volume - leaf.volume,
                           s.p_o - leaf.p_o + e,
                           s.volume_h};

    NodeStatL2<T> c_old = {leaf.volume,
                           leaf.p_o,
                           s.volume_h};

    return - merge(s_old, c_old, e);
}



template<typename T>
T Objectives<T>::lift(NodeStatL3<T> & s, NodeStatL3<T> & child, T e) {

    NodeStatL3<T> s_old = {s.volume - child.volume,
                           s.p_o - child.p_o + e,
                           s.volume_h,
                           s.p_a - child.p_a};

    NodeStatL3<T> c_old = {child.volume,
                           child.p_o,
                           s.volume_h,
                           child.p_a};

    return - merge(s_old, c_old, e);
}

template<typename T> T Objectives<T>::swap(NodeStatL3<T> &s, NodeStatL3<T> &t, NodeStatL3<T> &child, T e_sc, T e_ct) {
    NodeStatL3<T> s_old = {s.volume - child.volume,
                           s.p_o - child.p_o + e_sc,
                           s.volume_h,
                           s.p_a - child.p_a};
    NodeStatL3<T> c_old = {child.volume,
                           child.p_o,
                           s.volume_h,
                           child.p_a};
    return merge(c_old, t, e_ct) - merge(s_old, c_old, e_sc);

}

template<typename T>
T Objectives<T>::swap(NodeStatL2<T> &s, NodeStatL2<T> &t, NodeStatL2<T> &leaf, T e_sc, T e_ct) {
    NodeStatL2<T> s_old = {s.volume - leaf.volume,
                           s.p_o - leaf.p_o + e_sc,
                           s.volume_h};
    NodeStatL2<T> c_old = {leaf.volume,
                           leaf.p_o,
                           s.volume_h};
    return merge(c_old, t, e_ct) - merge(s_old, c_old, e_sc);
}

//template <typename T> T Objectives<T>::merge(T vol_a1, T vol_a2, T p_a1, T p_a2, T pi_a1, T pi_a2, T e, T v_gamma) {
//    return (pi_a1 - p_a1) * log2(vol_a1) + (pi_a2 - p_a2) * log2(vol_a2)
//        - (pi_a1 + pi_a2 - p_a1 - p_a2 - e * 2) * log2(vol_a1 + vol_a2)
//        - 2 * e * log2(v_gamma);
//}
//
//template <typename T> T Objectives<T>::merge(T vol_a1, T vol_a2, T p_a1, T p_a2, T e, T v_gamma) {
//    return (vol_a1 - p_a1) * log2(vol_a1) + (vol_a1 - p_a2) * log2(vol_a2)
//           - (pi_a1 + pi_a2 - p_a1 - p_a2 - e * 2) * log2(vol_a1 + vol_a2)
//           - 2 * e * log2(v_gamma);
//}

